import torch
import numpy


class CustomLoss(torch.nn.Module):
    def __init__(self):
        super(CustomLoss, self).__init__()

    def forward(self, predictions, targets):
        if targets.dim()==1:
            N = 1
        else:
            N = targets.shape[0]

        Wn = 1.0 - targets

        loss =  torch.tensordot(predictions, Wn)

        loss += torch.tensordot(1.0-predictions, targets)

        return loss / N